import requests
import json
import os
import time
from tqdm import tqdm
from api_request import query_by_paper_id

def save_paper_queries(paper_ids, output_dir="paper_queries", top_k=5):
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    exist_paper_id = set()
    for root, dirs, files in os.walk(output_dir):
        for file in files:
            if file.endswith('.json'):
                exist_paper_id.add(file.split('.')[0])
    
    print(f"已存在{len(exist_paper_id)}篇论文,共有{len(paper_ids)}篇论文")
    paper_ids = set(paper_ids) - exist_paper_id
    # 遍历所有paper_id并查询
    for paper_id in tqdm(paper_ids, desc="Querying papers"):
        try:
            # 查询论文
            result = query_by_paper_id(paper_id, top_k)
            
            # 构建输出文件路径
            output_file = os.path.join(output_dir, f"{paper_id}.json")
            
            # 保存结果
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(result, f, ensure_ascii=False, indent=4)
                
            # 添加短暂延迟避免请求过快
            # time.sleep(0.1)
            
        except Exception as e:
            print(f"Error processing paper_id {paper_id}: {str(e)}")
            continue

# 使用示例
if __name__ == "__main__":
    with open("paper_id_list.json",'r') as f:
        paper_ids = json.load(f)
    print(len(paper_ids))
    # 设置输出目录和每篇论文获取的片段数
    output_dir = "papers"
    top_k = 100
    
    # 执行查询和保存
    save_paper_queries(paper_ids, output_dir, top_k)
